import junit.framework.TestCase;

import java.util.ArrayList;
import java.util.Iterator;
import de.unibi.comet.hmm.*;

public class HMMTest extends TestCase {
	
	private static void assertConsistent(StateAndMassDistribution dist, int states) {
		double total_sum = 0.0;
		for (int state=0; state<states; ++state) {
			Iterator<Emission> iter = dist.getMassDistribution(state);
			double sum = 0.0;
			while (iter.hasNext()) {
				Emission e = iter.next();
				assertTrue(dist.get(state, e.getMass())==e.getProbability());
				sum+=e.getProbability();
			}
			assertEquals(sum, dist.getStateProbability(state), 1e-7);
			total_sum+=sum;
		}
		assertEquals(1.0, total_sum, 1e-7);
	}
	
	public void testStatistics() {
		ArrayList<Node> nodes = new ArrayList<Node>();
		ArrayList<Transition> t = new ArrayList<Transition>();
		t.add(new Transition(1, 0.6));
		t.add(new Transition(3, 0.4));
		nodes.add(new StandardNode(t));
		t = new ArrayList<Transition>();
		t.add(new Transition(2, 0.9));
		t.add(new Transition(3, 0.1));
		nodes.add(new StandardNode(t));
		ArrayList<Emission> e = new ArrayList<Emission>();
		e.add(new Emission(1, 1.0));
		t = new ArrayList<Transition>();
		t.add(new Transition(1, 1.0));
		nodes.add(new StandardNode(t,e));
		e = new ArrayList<Emission>();
		e.add(new Emission(0, 0.1));
		e.add(new Emission(5, 0.9));
		t = new ArrayList<Transition>();
		t.add(new Transition(2, 1.0));
		nodes.add(new StandardNode(t,e));
		
		HMM hmm = new HMM(nodes);
		StateAndMassDistribution dist = new ArrayStateAndMassDistribution(4,0,false);
		dist.set(0, 0, 1.0);
		assertConsistent(dist,4);
		assertEquals(1.0, dist.getStateProbability(0), 1e-7);
		
		dist = hmm.step(dist);
		assertConsistent(dist,4);
		assertEquals(0.0, dist.getStateProbability(0), 1e-7);
		assertEquals(0.0, dist.get(0,0), 1e-7);
		assertEquals(0.6, dist.getStateProbability(1), 1e-7);
		assertEquals(0.6, dist.get(1,0), 1e-7);
		assertEquals(0.0, dist.getStateProbability(2), 1e-7);
		assertEquals(0.0, dist.get(2,0), 1e-7);
		assertEquals(0.4, dist.getStateProbability(3), 1e-7);
		assertEquals(0.04, dist.get(3,0), 1e-7);
		assertEquals(0.36, dist.get(3,5), 1e-7);
		
		dist = hmm.step(dist);
		assertConsistent(dist,4);
		assertEquals(0.0, dist.getStateProbability(0), 1e-7);
		assertEquals(0.0, dist.get(0,0), 1e-7);
		assertEquals(0.0, dist.getStateProbability(1), 1e-7);
		assertEquals(0.0, dist.get(1,0), 1e-7);
		assertEquals(0.94, dist.getStateProbability(2), 1e-7);
		assertEquals(0.58, dist.get(2,1), 1e-7);
		assertEquals(0.36, dist.get(2,6), 1e-7);
		assertEquals(0.06, dist.getStateProbability(3), 1e-7);
		assertEquals(0.006, dist.get(3,0), 1e-7);
		assertEquals(0.054, dist.get(3,5), 1e-7);

		dist = hmm.step(dist);
		assertConsistent(dist,4);
		assertEquals(0.0, dist.getStateProbability(0), 1e-7);
		assertEquals(0.94, dist.getStateProbability(1), 1e-7);
		assertEquals(0.58, dist.get(1,1), 1e-7);
		assertEquals(0.36, dist.get(1,6), 1e-7);
		assertEquals(0.06, dist.getStateProbability(2), 1e-7);
		assertEquals(0.006, dist.get(2,1), 1e-7);
		assertEquals(0.054, dist.get(2,6), 1e-7);
		assertEquals(0.0, dist.getStateProbability(3), 1e-7);

		dist = hmm.step(dist);
		assertConsistent(dist,4);
		assertEquals(0.0, dist.getStateProbability(0), 1e-7);
		assertEquals(0.06, dist.getStateProbability(1), 1e-7);
		assertEquals(0.006, dist.get(1,1), 1e-7);
		assertEquals(0.054, dist.get(1,6), 1e-7);
		assertEquals(0.846, dist.getStateProbability(2), 1e-7);
		assertEquals(0.522, dist.get(2,2), 1e-7);
		assertEquals(0.324, dist.get(2,7), 1e-7);
		assertEquals(0.094, dist.getStateProbability(3), 1e-7);
		assertEquals(0.0058, dist.get(3,1), 1e-7);
		assertEquals(0.0558, dist.get(3,6), 1e-7);		
		assertEquals(0.0324, dist.get(3,11), 1e-7);		
	}
}
